from PyQt6.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QFormLayout, 
    QLineEdit, QComboBox, QSpinBox, QPushButton, 
    QLabel, QGroupBox, QMessageBox
)
from PyQt6.QtCore import Qt, pyqtSignal
from typing import Dict, Any, Optional


class ConnectionDialog(QDialog):
    """ChromaDB连接配置对话框"""
    
    # 自定义信号，用于传递连接配置
    connectionAccepted = pyqtSignal(dict)
    
    def __init__(self, parent=None, connection: Optional[Dict[str, Any]] = None):
        """
        初始化连接对话框
        
        Args:
            parent: 父窗口
            connection: 现有连接配置，用于编辑模式
        """
        super().__init__(parent)
        self.connection = connection
        self.setWindowTitle("ChromaDB连接配置")
        self.resize(500, 400)
        
        self._init_ui()
        
        # 如果提供了连接配置，则填充表单
        if connection:
            self._fill_form(connection)
    
    def _init_ui(self):
        """初始化UI组件"""
        main_layout = QVBoxLayout()
        
        # 基本连接信息
        basic_group = QGroupBox("基本连接信息")
        basic_layout = QFormLayout()
        
        self.name_edit = QLineEdit()
        basic_layout.addRow("连接名称:", self.name_edit)
        
        self.type_combo = QComboBox()
        self.type_combo.addItems(["local", "remote"])
        self.type_combo.currentTextChanged.connect(self._on_type_changed)
        basic_layout.addRow("连接类型:", self.type_combo)
        
        # 远程连接设置
        self.remote_group = QGroupBox("远程连接设置")
        remote_layout = QFormLayout()
        
        self.host_edit = QLineEdit()
        remote_layout.addRow("主机地址:", self.host_edit)
        
        self.port_spin = QSpinBox()
        self.port_spin.setRange(1, 65535)
        self.port_spin.setValue(8000)
        remote_layout.addRow("端口:", self.port_spin)
        
        self.remote_group.setLayout(remote_layout)
        
        # 本地连接设置
        self.local_group = QGroupBox("本地连接设置")
        local_layout = QFormLayout()
        
        self.path_edit = QLineEdit()
        local_layout.addRow("数据路径:", self.path_edit)
        
        self.local_group.setLayout(local_layout)
        
        # 嵌入模型设置
        embedding_group = QGroupBox("嵌入模型设置")
        embedding_layout = QFormLayout()
        
        self.embedding_type_combo = QComboBox()
        self.embedding_type_combo.addItems(["local", "openai"])
        self.embedding_type_combo.currentTextChanged.connect(self._on_embedding_type_changed)
        embedding_layout.addRow("嵌入类型:", self.embedding_type_combo)
        
        self.embedding_model_edit = QLineEdit()
        embedding_layout.addRow("嵌入模型:", self.embedding_model_edit)
        
        # OpenAI设置
        self.openai_group = QGroupBox("OpenAI设置")
        openai_layout = QFormLayout()
        
        self.base_url_edit = QLineEdit()
        self.base_url_edit.setText("https://api.openai.com/v1")
        openai_layout.addRow("API基础URL:", self.base_url_edit)
        
        self.api_key_edit = QLineEdit()
        self.api_key_edit.setEchoMode(QLineEdit.EchoMode.Password)
        openai_layout.addRow("API密钥:", self.api_key_edit)
        
        self.openai_group.setLayout(openai_layout)
        
        # 按钮
        button_layout = QHBoxLayout()
        self.cancel_button = QPushButton("取消")
        self.cancel_button.clicked.connect(self.reject)
        
        self.save_button = QPushButton("保存")
        self.save_button.clicked.connect(self._on_save)
        self.save_button.setDefault(True)
        
        button_layout.addWidget(self.cancel_button)
        button_layout.addWidget(self.save_button)
        
        # 设置布局
        basic_group.setLayout(basic_layout)
        embedding_group.setLayout(embedding_layout)
        
        main_layout.addWidget(basic_group)
        main_layout.addWidget(self.remote_group)
        main_layout.addWidget(self.local_group)
        main_layout.addWidget(embedding_group)
        main_layout.addWidget(self.openai_group)
        main_layout.addLayout(button_layout)
        
        self.setLayout(main_layout)
        
        # 初始化UI状态
        self._on_type_changed(self.type_combo.currentText())
        self._on_embedding_type_changed(self.embedding_type_combo.currentText())
    
    def _on_type_changed(self, type_text: str):
        """
        连接类型变更处理
        
        Args:
            type_text: 连接类型文本
        """
        if type_text == "local":
            self.remote_group.setVisible(False)
            self.local_group.setVisible(True)
        else:
            self.remote_group.setVisible(True)
            self.local_group.setVisible(False)
    
    def _on_embedding_type_changed(self, type_text: str):
        """
        嵌入类型变更处理
        
        Args:
            type_text: 嵌入类型文本
        """
        if type_text == "local":
            self.openai_group.setVisible(False)
        else:
            self.openai_group.setVisible(True)
    
    def _fill_form(self, connection: Dict[str, Any]):
        """
        填充表单数据
        
        Args:
            connection: 连接配置
        """
        self.name_edit.setText(connection.get('name', ''))
        
        conn_type = connection.get('type', 'remote')
        self.type_combo.setCurrentText(conn_type)
        
        if conn_type == 'remote':
            self.host_edit.setText(connection.get('host', ''))
            self.port_spin.setValue(connection.get('port', 8000))
        else:
            self.path_edit.setText(connection.get('path', ''))
        
        embedding_type = connection.get('embedding_type', 'local')
        self.embedding_type_combo.setCurrentText(embedding_type)
        self.embedding_model_edit.setText(connection.get('embedding_model', ''))
        
        if embedding_type == 'openai':
            self.base_url_edit.setText(connection.get('base_url', 'https://api.openai.com/v1'))
            self.api_key_edit.setText(connection.get('api_key', ''))
    
    def _on_save(self):
        """保存按钮点击处理"""
        # 验证必填字段
        if not self.name_edit.text().strip():
            QMessageBox.warning(self, "验证错误", "连接名称不能为空")
            return
        
        # 构建连接配置
        connection = {
            'name': self.name_edit.text().strip(),
            'type': self.type_combo.currentText(),
            'embedding_type': self.embedding_type_combo.currentText(),
            'embedding_model': self.embedding_model_edit.text().strip()
        }
        
        # 根据连接类型添加字段
        if connection['type'] == 'remote':
            connection['host'] = self.host_edit.text().strip()
            connection['port'] = self.port_spin.value()
            
            if not connection['host']:
                QMessageBox.warning(self, "验证错误", "主机地址不能为空")
                return
        else:
            connection['path'] = self.path_edit.text().strip()
            
            if not connection['path']:
                QMessageBox.warning(self, "验证错误", "数据路径不能为空")
                return
        
        # 根据嵌入类型添加字段
        if connection['embedding_type'] == 'openai':
            connection['base_url'] = self.base_url_edit.text().strip()
            connection['api_key'] = self.api_key_edit.text().strip()
            
            if not connection['base_url'] or not connection['api_key']:
                QMessageBox.warning(self, "验证错误", "OpenAI API基础URL和密钥不能为空")
                return
        
        # 验证嵌入模型
        if not connection['embedding_model']:
            QMessageBox.warning(self, "验证错误", "嵌入模型不能为空")
            return
        
        # 发送信号并关闭对话框
        self.connectionAccepted.emit(connection)
        self.accept()
